import openai
import os


class ChatBot:
    """
    ChatBot 类用于管理与 OpenAI API 的交互和对话历史记录。

    Attributes:
        client (openai.Client): OpenAI API 客户端。
        history (list): 存储对话历史的列表。
    """

    def __init__(self, api_key):
        """
        初始化 ChatBot 实例。

        Args:
            api_key (str): OpenAI API 密钥。
        """
        self.client = openai.Client(
            base_url="https://api.xverse.cn/v1",
            api_key=api_key
        )
        self.history = []

    def add_message(self, role, content):
        """
        将消息添加到对话历史中。

        Args:
            role (str): 消息角色，可以是 "user" 或 "assistant"。
            content (str): 消息内容。
        """
        self.history.append({"role": role, "content": content})

    def get_response(self, user_input):
        """
        获取用户输入的响应。

        Args:
            user_input (str): 用户输入的消息。
        """
        # 将用户输入添加到对话历史中
        self.add_message("user", user_input)

        # 调用 OpenAI API 获取响应
        result = self.client.chat.completions.create(
            model="XVERSE-13B-2",
            messages=self.history,
            stream=True
        )

        # 初始化响应内容
        response_content = ""

        # 处理流式响应
        for chunk in result:
            if chunk.choices[0].delta.content is not None:
                response_content += chunk.choices[0].delta.content
                print(chunk.choices[0].delta.content, end='', flush=True)

        # 将助手的响应添加到对话历史中
        self.add_message("assistant", response_content)


if __name__ == "__main__":
    # 获取环境变量中的 API 密钥
    api_key = os.getenv("XVERSE_API_KEY")

    # 创建 ChatBot 实例
    chatbot = ChatBot(api_key)

    # 进入对话循环
    while True:
        # 获取用户输入
        user_input = input("You: ")

        # 如果用户输入 "exit" 或 "quit"，退出循环
        if user_input.lower() in ["exit", "quit"]:
            break

        # 获取响应并打印
        chatbot.get_response(user_input)
        print()  # 换行以分隔对话
